import os

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from dotenv import load_dotenv
from openai import OpenAI as DeepSeekClient
from dashscope import Generation as QwenClient

# 加载环境变量
load_dotenv()

app = FastAPI()

# 配置模型参数
MODEL_CONFIG = {
    "deepseek": {
        "client": DeepSeekClient(
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            base_url=os.getenv("DEEPSEEK_BASE_URL")
        ),
        "model_name": "deepseek-chat"
    },
    "qwen": {
        "client": QwenClient,
        "model_name": "qwen-plus"
    }
}


# 统一请求格式（兼容OpenAI）
class ChatRequest(BaseModel):
    model: str  # deepseek 或 qwen
    messages: list
    temperature: float = 0.7
    max_tokens: int = 1024


# 统一响应格式
class ChatResponse(BaseModel):
    model: str
    content: str


@app.post("/v1/chat")
async def chat_completion(request: ChatRequest):
    model_type = request.model.lower()
    if model_type not in MODEL_CONFIG:
        raise HTTPException(400, f"Unsupported model: {model_type}")

    try:
        if model_type == "deepseek":
            response = MODEL_CONFIG["deepseek"]["client"].chat.completions.create(
                model=MODEL_CONFIG["deepseek"]["model_name"],
                messages=request.messages,
                temperature=request.temperature,
                max_tokens=request.max_tokens
            )
            content = response.choices[0].message.content

        elif model_type == "qwen":
            response = QwenClient.call(
                model=MODEL_CONFIG["qwen"]["model_name"],
                messages=request.messages,
                temperature=request.temperature,
                max_tokens=request.max_tokens,
                result_format='message'
            )
            content = response.output.choices[0].message.content

        return ChatResponse(model=model_type, content=content)

    except Exception as e:
        raise HTTPException(500, f"API Error: {str(e)}")


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8090)
